# -*- coding: utf-8 -*-

"""
Cont.

Usage:
  cont.py [-e <e_value> -n <aligns_number> -t <reads_number> ] -f <fastq_file>
  cont.py [-e <e_value> -n <aligns_number> -a] -f <fastq_file>
  cont.py [-e <e_value> -n <aligns_number> (-g (-a | -t <reads_number>))] -f <fastq_file>
  cont.py [(-G (-a | -t <reads_number>))] -f <fastq_file>

Options:
  -h --help             Show this screen.
  --version             Show version.
  -e <e_value>          E-value threshold [default: 0.5].
  -n <aligns_number>    Number of alignments stored for each read [default: 1].
  -t <reads_number>     Number of reads to process from file.
  -a                    Process all reads from file.
  -g                    Plot GC% 
  -G                    Plot GC% only (no read BLASTing)
  -f                    FASTQ file name
"""

import sys

modules = ["sys", "os", "scipy", "math", "matplotlib", "docopt", "Bio"]
exit_flag = False
for module in modules:
    try:
        __import__(module)
    except ImportError:
        exit_flag = True
        sys.stderr.write("Error: Python module " + module + " is not installed.\n")
        
if exit_flag:
    sys.exit("Some required Python modules are not installed.")
        
import os
import scipy.stats as stats
import math
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from docopt import docopt
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML
from Bio import SeqIO
from Bio import SeqUtils


def blast_read(record, blast_result_filename):
    sys.stderr.write("BLASTing read " + record.id + ":\n")
    sys.stderr.write("Using BLAST on NCBI website...")
    try:
        result_handle = NCBIWWW.qblast("blastn", "nt", record.seq, None, 
                                   None, None, None, '(none)', 10.0, 
                                   None, None, None, 50, None, None, 
                                   None, None, None, None, None, None, 
                                   None, None, None, query_from=1, 
                                   query_to=100)
    except:
        sys.stderr.write("BLASTed\n")
    sys.stderr.write("BLASTed\n")
    with open(blast_result_filename, "w") as save_file:
        save_file.write(result_handle.read())
    

def parse_blast_record(blast_result_filename, e_limit, max_n, n, record_id):
    with open(blast_result_filename, "r") as result_handle:
        blast_record = NCBIXML.read(result_handle)
        align_exists = False
        titles = []
        sequences = []
        aligns = []
        for num, alignment in enumerate(blast_record.alignments[:n], 1):
            e_value = alignment.hsps[0].expect
            al = []
            al.append(str(record_id) + '\n') 
            if e_value <= e_limit:
                align_exists = True
                al.append('> Alignment ' + str(num) + ':\n')
                al.append('\n')
                al.append('Sequence title:\t' + alignment.title + '\n')
                al.append('Sequence length:\t' + str(alignment.length) + ' bp\n')
                al.append('Sequence e value:\t' + str(e_value) + '\n')
                al.append('\n')
                aligns.append(al)
            
                if single_read and num <= max_n:
                    align = alignment.hsps[0].sbjct
                    if len(align) != len(record.seq):
                        align += " " * (len(record.seq) - 
                                        len(alignment.hsps[0].sbjct))
                    sequences.append(align)
                    titles.append(alignment.title)

        if align_exists == False:
            al = []
            al.append('No alignments with e < ' + str(e) + '\n')
            aligns.append(al)
        
    return aligns, sequences, titles


def output_aligns(aligns_filename, aligns):
    with open(aligns_filename, 'w') as out:
        for al in aligns:
            for line in al:
                out.write(line)


def add_gc_content(gc_table, sequence):
    gc_content = SeqUtils.GC(sequence)
    gc_content_frac = gc_content - int(gc_content)
    if gc_content_frac >= 0.5:
        gc_content = int(math.ceil(gc_content))
    else:
        gc_content = int(gc_content)
                
    if gc_content not in gc_table:
        gc_table[gc_content] = 0
    gc_table[gc_content] += 1
    
    return gc_table
        
        
def calc_th_distr(gc_table):
    sys.stderr.write("Calculating theoretical GC%...\n")
    first_mode = 0
    max_num = 0
    
    reads_count = sum(gc_table.itervalues())
    first_mode = max(gc_table.iteritems(), key=(lambda i: i[1]))[0]
    max_num = max(first_mode, max_num)
            
    modes_sum = 0
    mode_dup_count = 0
    for i in range(first_mode, 100) + range(first_mode - 1, 0, -1):
        if i in gc_table.keys() and gc_table[i] > 0.95 * gc_table[first_mode]:
            modes_sum += i
            mode_dup_count +=1
        else:
            break

    if mode_dup_count != 0:
        mode = modes_sum / mode_dup_count
    else:
        mode = modes_sum
               
    st_deviation = 0                
    for i in gc_table.keys():
        st_deviation += ((i - mode)**2) * gc_table[i]
    st_deviation /= reads_count - 1
    st_deviation = math.sqrt(st_deviation)
    
    th_distr = []
    for i in range(101):
        prob = stats.norm.pdf([i], mode, st_deviation)
        th_distr.append(prob * reads_count)
        if (th_distr[i] > max_num): 
            max_num = th_distr[i]
            
    return (th_distr, max_num)


def plot_gc_distribs(max_num, gc_content_keys, gc_content_vals, th_distr):
    plt.xlabel(u'%GC')
    plt.ylabel(u'Number of reads')
    plt.title(u'%GC (' + filename + ')')
   
    max_num = max(max_num, max(gc_content_vals))
    yscale = 1.05
    plt.plot(gc_content_keys, gc_content_vals, label = 'Real GC%')
    plt.plot(range(101), th_distr, label = 'Theoretical GC%')
    plt.legend(loc='upper left', shadow=False)
    plt.axis([0, 100, 0, max_num * yscale])
    plt.savefig(gc_plot_filename)
    

def plot_quality_align(record, n, max_n, sequences, titles):
    sys.stderr.write("Plotting quality bar and table of alignments...\n")
    if n > max_n:
        sys.stderr.write("Warning: only first " + str(max_n) + " alignments will be included in the table.\n")
        
    bar_list = record.letter_annotations["phred_quality"]
    columns = tuple(''.join(record.seq))
    data = sequences
    n_rows = len(data)

    h_size = 8
    v_size = 4
    fig = plt.figure(figsize = (h_size, v_size), dpi = 200)
    
    delta1 = 0.02
    delta2 = 0.06
    delta3 = 0.02
    row_height = delta2 * (1 - 0.02 * (n_rows - 1))    
    x1 = 0.07
    y1 = ((delta1 + math.ceil(n_rows / 2.0) * delta2 / 2.0 + delta3) + 
         (row_height * n_rows))
    x2 = 0.90
    y2 = 0.90 - y1
    axes = plt.axes([x1, y1, x2, y2])
    axes.bar(range(len(record.seq)), bar_list, width = 1)
    
    axes.set_ylabel("Quality")
    axes.set_xticklabels([])
    axes.set_title('Quality and ' + str(n) + ' best alignments for read ' + 
                    str(record.id), fontsize=10)
    cell_text = []
    for row in data:
        row_text = []
        for i in range(len(row)):
            row_text.append(str(row[i]))
        cell_text.append(row_text)

    cell_colors = []
    row_labels = []
    for j in range(len(data)):
        h_colors=[]
        row_labels.append(str(j + 1))
        for i in range(len(columns)):
            color = '#00ee00' if data[j][i] == record.seq[i] else '#ee0000'
            h_colors.append(matplotlib.colors.colorConverter.to_rgb(color))
        cell_colors.append(h_colors)
                
    align_table = plt.table(cellText=cell_text, 
                          cellColours=cell_colors,
                          cellLoc='center', colWidths=None,
                          rowLabels=row_labels, rowColours=None, 
                          rowLoc='center',
                          colLabels=columns, colColours=None, 
                          colLoc='center',
                          loc='bottom', bbox=None)

    fig.subplots_adjust()
    
    textstr = ""
    for num, title in enumerate(titles, 1):
        if num > max_n or num > n:
            break
        if num % 2 != 0 and num != 1:
            textstr += '\n'
        elif num % 2 == 0:
            textstr += '$\,$'
        textstr += '%d:$\,$%s'%(num, title.split('|')[4])

    delta4 = 0.04
    text_axes = plt.axes([x1 - delta4, delta1, x2, y1 - n_rows * row_height])
    text_axes.axis('off')
    plt.text(0, 0, textstr, transform=text_axes.transAxes, fontsize=8)
    fig.savefig(align_plot_filename, dpi = 200)            

def is_non_zero_file(fpath):  
    return True if os.path.isfile(fpath) and os.path.getsize(fpath) > 0 else False


if __name__ == '__main__':
    arguments = docopt(__doc__, version='Cont 0.1')
    try:
        e = float(arguments["-e"])
        if (e <= 0):
            raise Exception
    except:
        sys.exit("Error: e-value (-e option) must be a positive floating point number.")
        
    try:
        n = int(arguments["-n"])
        if ((not arguments["-n"].isdigit()) or n <= 0):
            raise Exception
    except:
        sys.exit("Error: number of results (-n option) must be a positive integer.")
        
    a = arguments["-a"]
    g = arguments["-g"]
    G = arguments["-G"]
    t = 1
    if arguments["-t"] != None:
        try:
            t = int(arguments["-t"])
            if ((not arguments["-t"].isdigit()) or t <= 0):
                raise Exception
        except:
            sys.exit("Error: number of reads (-t option) must be a positive integer.")

filename = arguments["<fastq_file>"]
if (not is_non_zero_file(filename)):
    sys.exit("Error: file " + filename + " is empty or doesn't exist.")

gc_plot_filename = 'gc_plot.png'
align_plot_filename = 'quality_align_plot.png'
aligns_filename = 'aligns.txt'
blast_result_filename = 'blast_result.xml'

with open(filename, 'r') as handle:
    sys.stderr.write("Processing started...\n")
    single_read = False
    if arguments["-t"] == None and a == False:
        single_read = True
    
    max_n = 10
    gc_table = {}
    read_counter = 1
    first = True
    for record in SeqIO.parse(handle, "fastq"):
        if G == False:
            blast_read(record, blast_result_filename)
            aligns, sequences, titles = parse_blast_record(blast_result_filename,
                                                           e, max_n, n, 
                                                           record.id)
            output_aligns(aligns_filename, aligns)

            if g == True:
                gc_table = add_gc_content(gc_table, record.seq)   
        else:
            if first == True:
                first = False
                sys.stderr.write("Calculating real GC%...\n")
            gc_table = add_gc_content(gc_table, record.seq)
        
        if a == False:
            read_counter += 1
            if read_counter > t:
                break

    if g == True or G == True:
        th_distr, max_num = calc_th_distr(gc_table)
        gc_content_keys = []
        gc_content_vals = []
        for key in sorted(gc_table.keys()):
            gc_content_keys.append(key)
            gc_content_vals.append(gc_table[key])
        plot_gc_distribs(max_num, gc_content_keys, gc_content_vals, th_distr)

    if single_read:
        plot_quality_align(record, n, max_n, sequences, titles)
    
    sys.stderr.write("Processing finished!\n")
